package com.abc.util;

import com.abc.config.ZhipuConfig;
import com.abc.dto.request.TokenCountRequest;
import com.abc.dto.response.TokenCountResponse;
import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.*;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

import java.util.HashMap;
import java.util.Map;

/**
 * Token计算工具类
 * 对接智谱AI tokenizer API，用于计算文本的tokens数
 */
@Slf4j
@Component
public class TokenUtil {

    private final RestTemplate restTemplate;
    private final ZhipuConfig zhipuConfig;

    public TokenUtil(RestTemplate restTemplate, ZhipuConfig zhipuConfig) {
        this.restTemplate = restTemplate;
        this.zhipuConfig = zhipuConfig;
    }

    /**
     * 计算文本的token数量
     *
     * @param text 要计算token的文本
     * @return token数量
     */
    public Integer countTokens(String text) {
        return countTokens(text, "glm-4-flash");
    }

    /**
     * 计算文本的token数量
     *
     * @param text  要计算token的文本
     * @param model 模型名称
     * @return token数量
     */
    public Integer countTokens(String text, String model) {
        TokenCountRequest request = new TokenCountRequest();
        request.setText(text);
        request.setModel(model);
        
        TokenCountResponse response = countTokensWithDetails(request);
        return response != null && response.getUsage() != null ? response.getUsage().getPromptTokens() : 0;
    }

    /**
     * 计算文本的token数量（详细响应）
     *
     * @param request token计算请求
     * @return 详细的token计算响应
     */
    public TokenCountResponse countTokensWithDetails(TokenCountRequest request) {
        try {
            // 构建请求体
            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", request.getModel());
            
            Map<String, Object> message = new HashMap<>();
            message.put("role", request.getRole());
            message.put("content", request.getText());
            requestBody.put("messages", new Object[]{message});

            // 设置请求头
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            headers.setBearerAuth(zhipuConfig.getApiKey());

            // 创建请求实体
            HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);

            // 发送请求
            ResponseEntity<String> response = restTemplate.exchange(
                    zhipuConfig.getTokenizerUrl(),
                    HttpMethod.POST,
                    entity,
                    String.class
            );

            if (response.getStatusCode() == HttpStatus.OK) {
                log.debug("Token计算成功，响应: {}", response.getBody());
                return JSON.parseObject(response.getBody(), TokenCountResponse.class);
            } else {
                log.error("Token计算失败，状态码: {}, 响应: {}", response.getStatusCode(), response.getBody());
                return null;
            }

        } catch (Exception e) {
            log.error("Token计算异常，文本: {}", request.getText(), e);
            return null;
        }
    }

    /**
     * 批量计算多个文本的token数量
     *
     * @param texts 文本数组
     * @return 总token数量
     */
    public Integer countTokensBatch(String[] texts) {
        return countTokensBatch(texts, "glm-4-flash");
    }

    /**
     * 批量计算多个文本的token数量
     *
     * @param texts 文本数组
     * @param model 模型名称
     * @return 总token数量
     */
    public Integer countTokensBatch(String[] texts, String model) {
        int totalTokens = 0;
        for (String text : texts) {
            totalTokens += countTokens(text, model);
        }
        return totalTokens;
    }

    /**
     * 检查文本是否超过指定的token限制
     *
     * @param text        要检查的文本
     * @param maxTokens   最大token限制
     * @return 是否超过限制
     */
    public boolean isTextExceedsLimit(String text, int maxTokens) {
        return isTextExceedsLimit(text, maxTokens, "glm-4-flash");
    }

    /**
     * 检查文本是否超过指定的token限制
     *
     * @param text        要检查的文本
     * @param maxTokens   最大token限制
     * @param model       模型名称
     * @return 是否超过限制
     */
    public boolean isTextExceedsLimit(String text, int maxTokens, String model) {
        Integer tokenCount = countTokens(text, model);
        return tokenCount != null && tokenCount > maxTokens;
    }

    /**
     * 截断文本到指定的token限制
     *
     * @param text        要截断的文本
     * @param maxTokens   最大token限制
     * @return 截断后的文本
     */
    public String truncateTextToTokenLimit(String text, int maxTokens) {
        return truncateTextToTokenLimit(text, maxTokens, "glm-4-flash");
    }

    /**
     * 截断文本到指定的token限制
     * 注意：这是一个简单的实现，实际应用中可能需要更复杂的截断策略
     *
     * @param text        要截断的文本
     * @param maxTokens   最大token限制
     * @param model       模型名称
     * @return 截断后的文本
     */
    public String truncateTextToTokenLimit(String text, int maxTokens, String model) {
        if (text == null || text.isEmpty()) {
            return text;
        }

        // 如果文本token数不超过限制，直接返回
        if (!isTextExceedsLimit(text, maxTokens, model)) {
            return text;
        }

        // 简单的截断策略：按字符数比例截断
        // 注意：这不是最准确的方法，实际应用中建议使用更精确的截断算法
        int currentTokens = countTokens(text, model);
        if (currentTokens == 0) {
            return text;
        }

        double ratio = (double) maxTokens / currentTokens;
        int targetLength = (int) (text.length() * ratio);
        
        // 确保不会截断到空字符串
        targetLength = Math.max(1, targetLength);
        
        return text.substring(0, Math.min(targetLength, text.length()));
    }
}
